import time
import os
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from PIL import Image
import paddleocr
import re
import cv2
import numpy as np

def setup_driver():
    """设置Chrome浏览器选项"""
    chrome_options = Options()
    chrome_options.add_argument('--disable-gpu')
    chrome_options.add_argument('--no-sandbox')
    chrome_options.add_argument('--disable-dev-shm-usage')
    chrome_options.add_argument('--window-size=375,812')  # 模拟手机屏幕尺寸
    chrome_options.add_argument('--user-agent=Mozilla/5.0 (iPhone; CPU iPhone OS 14_0 like Mac OS X) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/14.0 Mobile/15E148 Safari/604.1')
    
    driver = webdriver.Chrome(options=chrome_options)
    return driver

def take_screenshot(driver, url, filename):
    """访问URL并截图"""
    print(f"正在访问: {url}")
    driver.get(url)
    time.sleep(5)  # 等待页面加载
    
    # 截图并保存
    driver.save_screenshot(filename)
    print(f"截图已保存: {filename}")
    return filename

def scroll_and_screenshot(driver, scroll_pixels, filename):
    """滚动页面并截图"""
    driver.execute_script(f"window.scrollBy(0, {scroll_pixels});")
    time.sleep(3)  # 等待滚动完成
    driver.save_screenshot(filename)
    print(f"滚动后截图已保存: {filename}")
    return filename

def preprocess_image(image_path):
    """预处理图像以提高OCR识别率"""
    # 读取图像
    image = cv2.imread(image_path)
    # 转换为灰度图
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    # 应用自适应阈值处理
    processed = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
                                     cv2.THRESH_BINARY, 11, 2)
    return processed

def extract_prices_with_paddleocr(image_path):
    """使用PaddleOCR从图片中提取价格信息"""
    try:
        # 初始化PaddleOCR
        ocr = paddleocr.PaddleOCR(use_angle_cls=True, lang='ch')
        
        # 预处理图像
        processed_image = preprocess_image(image_path)
        
        # 使用PaddleOCR进行识别
        result = ocr.ocr(processed_image, cls=True)
        
        # 提取所有识别到的文本
        all_text = []
        for line in result:
            if line and len(line) > 0:
                for word_info in line:
                    if word_info and len(word_info) > 1:
                        text = word_info[1][0]
                        confidence = word_info[1][1]
                        if confidence > 0.5:  # 只保留置信度较高的结果
                            all_text.append(text)
        
        # 使用正则表达式查找价格模式
        price_pattern = r'[¥￥]?\s*\d+\.?\d*\s*[元]?'
        prices = []
        for text in all_text:
            found_prices = re.findall(price_pattern, text)
            prices.extend(found_prices)
        
        return prices
    except Exception as e:
        print(f"使用PaddleOCR处理图片时出错: {e}")
        return []

def main():
    # 目标URL
    url = "https://mobile.pinduoduo.com/search_result.html?search_key=%E6%84%9F%E5%86%92%E7%81%B5&search_type=goods&source=index&options=3&search_met_track=manual&refer_page_el_sn=99884&refer_page_name=search_result&refer_page_id=10015_1757120538005_0pkr00u060&refer_page_sn=10015"
    
    # 创建截图目录
    screenshot_dir = "screenshots"
    os.makedirs(screenshot_dir, exist_ok=True)
    
    # 设置浏览器驱动
    driver = setup_driver()
    
    try:
        # 初始截图
        screenshot1 = os.path.join(screenshot_dir, "screenshot1.png")
        take_screenshot(driver, url, screenshot1)
        
        # 滚动页面并再次截图
        screenshot2 = os.path.join(screenshot_dir, "screenshot2.png")
        scroll_and_screenshot(driver, 1000, screenshot2)
        
        # 再次滚动并截图
        screenshot3 = os.path.join(screenshot_dir, "screenshot3.png")
        scroll_and_screenshot(driver, 1000, screenshot3)
        
        # 处理所有截图并提取价格
        screenshot_files = [screenshot1, screenshot2, screenshot3]
        
        for i, screenshot_file in enumerate(screenshot_files, 1):
            print(f"\n处理截图 {i}: {screenshot_file}")
            prices = extract_prices_with_paddleocr(screenshot_file)
            
            if prices:
                print(f"找到 {len(prices)} 个价格:")
                for price in prices:
                    print(f"  - {price}")
            else:
                print("未找到价格信息")
                
    except Exception as e:
        print(f"程序执行出错: {e}")
    finally:
        driver.quit()
        print("\n浏览器已关闭")

if __name__ == "__main__":
    # 检查PaddleOCR是否可用
    try:
        import paddleocr
        print("PaddleOCR已安装，开始执行程序...")
    except ImportError:
        print("请先安装PaddleOCR: pip install paddlepaddle paddleocr")
        exit(1)
        
    main()